package org.loong.crypto.core.utils;

import java.math.BigInteger;
import java.security.KeyFactory;
import java.security.NoSuchAlgorithmException;
import java.security.PrivateKey;
import java.security.Provider;
import java.security.PublicKey;
import java.security.interfaces.RSAPrivateKey;
import java.security.interfaces.RSAPublicKey;
import java.security.spec.InvalidKeySpecException;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.RSAPrivateCrtKeySpec;
import java.security.spec.RSAPublicKeySpec;
import java.security.spec.X509EncodedKeySpec;

import org.bouncycastle.util.BigIntegers;
import org.loong.crypto.common.exception.CryptoException;
import org.loong.crypto.common.exception.InvalidKeyException;
import org.loong.crypto.core.GlobalBouncyCastleProvider;

/**
 * RSA Key conversion utility.
 */
public class RSAKeyUtils {

    /**
     * Returns the RSA public key of the specified RSA key.
     *
     * @param key the RSA key bytes. Must not be {@code null}.
     * @return the RSA public key.
     * @throws CryptoException the RSA key transformation failed.
     */
    public static RSAPublicKey toRSAPublicKey(final byte[] key) throws CryptoException {
        final Provider provider = GlobalBouncyCastleProvider.INSTANCE.getProvider();
        try {
            KeyFactory keyFactory = KeyFactory.getInstance("RSA", provider);
            X509EncodedKeySpec keySpec = new X509EncodedKeySpec(key);
            RSAPublicKey publicKey = (RSAPublicKey) keyFactory.generatePublic(keySpec);
            if (publicKey == null) {
                throw new InvalidKeyException("The RSA public key transformation failed");
            }
            return publicKey;
        } catch (NoSuchAlgorithmException | InvalidKeySpecException e) {
            throw new InvalidKeyException("The RSA public key transformation failed");
        }
    }

    /**
     * Returns the RSA public key.
     *
     * @param modulus the modulus.
     * @param publicExponent the public exponent.
     * @return the RSA public key.
     */
    public static RSAPublicKey convertRSAPublicKey(byte[] modulus, byte[] publicExponent) {
        final Provider provider = GlobalBouncyCastleProvider.INSTANCE.getProvider();
        try {
            final RSAPublicKeySpec keySpec = new RSAPublicKeySpec(BigIntegers.fromUnsignedByteArray(modulus), BigIntegers.fromUnsignedByteArray(publicExponent));
            KeyFactory keyFactory = KeyFactory.getInstance("RSA", provider);
            RSAPublicKey publicKey = (RSAPublicKey) keyFactory.generatePublic(keySpec);
            if (publicKey == null) {
                throw new InvalidKeyException("The RSA public key transformation failed");
            }
            return publicKey;
        } catch (NoSuchAlgorithmException | InvalidKeySpecException e) {
            throw new InvalidKeyException("The RSA public key transformation failed");
        }
    }

    /**
     * Returns the RSA private key of the specified RSA key.
     *
     * @param key the RSA key bytes. Must not be {@code null}.
     * @return the RSA private key.
     * @throws CryptoException the RSA key transformation failed.
     */
    public static RSAPrivateKey toRSAPrivateKey(final byte[] key) throws CryptoException {
        final Provider provider = GlobalBouncyCastleProvider.INSTANCE.getProvider();
        try {
            KeyFactory keyFactory = KeyFactory.getInstance("RSA", provider);
            PKCS8EncodedKeySpec keySpec = new PKCS8EncodedKeySpec(key);
            RSAPrivateKey privateKey = (RSAPrivateKey) keyFactory.generatePrivate(keySpec);
            if (privateKey == null) {
                throw new InvalidKeyException("The RSA private key transformation failed");
            }
            return privateKey;
        } catch (NoSuchAlgorithmException | InvalidKeySpecException e) {
            throw new InvalidKeyException("The RSA private key transformation failed");
        }
    }

    /**
     * Returns the RSA private key.
     *
     * @param modulus the modulus n
     * @param publicExponent the public exponent e
     * @param privateExponent the private exponent d
     * @param primeP the prime factor p of n
     * @param primeQ the prime factor q of n
     * @param primeExponentP this is d mod (p-1)
     * @param primeExponentQ this is d mod (q-1)
     * @param crtCoefficient the Chinese Remainder Theorem
     * @return the RSA public key.
     */
    public static RSAPrivateKey convertRSAPrivateKey(byte[] modulus, byte[] publicExponent, byte[] privateExponent, byte[] primeP, byte[] primeQ, byte[] primeExponentP,
            byte[] primeExponentQ, byte[] crtCoefficient) {
        final Provider provider = GlobalBouncyCastleProvider.INSTANCE.getProvider();
        try {
            BigInteger m = BigIntegers.fromUnsignedByteArray(modulus);
            BigInteger e = BigIntegers.fromUnsignedByteArray(publicExponent);
            BigInteger d = BigIntegers.fromUnsignedByteArray(privateExponent);
            BigInteger p = BigIntegers.fromUnsignedByteArray(primeP);
            BigInteger q = BigIntegers.fromUnsignedByteArray(primeQ);
            BigInteger dp = BigIntegers.fromUnsignedByteArray(primeExponentP);
            BigInteger dq = BigIntegers.fromUnsignedByteArray(primeExponentQ);
            BigInteger qi = BigIntegers.fromUnsignedByteArray(crtCoefficient);

            final RSAPrivateCrtKeySpec keySpec = new RSAPrivateCrtKeySpec(m, e, d, p, q, dp, dq, qi);
            KeyFactory keyFactory = KeyFactory.getInstance("RSA", provider);
            RSAPrivateKey privateKey = (RSAPrivateKey) keyFactory.generatePrivate(keySpec);
            if (privateKey == null) {
                throw new InvalidKeyException("The RSA private key transformation failed");
            }
            return privateKey;
        } catch (NoSuchAlgorithmException | InvalidKeySpecException e) {
            throw new InvalidKeyException("The RSA public key transformation failed");
        }
    }

    /**
     * Returns the length in bits of the specified RSA public key.
     *
     * @param publicKey The RSA public key. Must not be {@code null}.
     * @return The key length in bits.
     */
    public static int keyBitLength(final PublicKey publicKey) {
        if (!(publicKey instanceof RSAPublicKey)) {
            return -1; // May be an PKCS#11 backed key
        }

        RSAPublicKey rsaPublicKey = (RSAPublicKey) publicKey;
        try {
            return rsaPublicKey.getModulus().bitLength();
        } catch (Exception e) {
            return -1;
        }
    }

    /**
     * Returns the length in bits of the specified RSA private key.
     *
     * @param privateKey The RSA private key. Must not be {@code null}.
     * @return The key length in bits.
     */
    public static int keyBitLength(final PrivateKey privateKey) {
        if (!(privateKey instanceof RSAPrivateKey)) {
            return -1; // May be an PKCS#11 backed key
        }

        RSAPrivateKey rsaPrivateKey = (RSAPrivateKey) privateKey;
        try {
            return rsaPrivateKey.getModulus().bitLength();
        } catch (Exception e) {
            return -1;
        }
    }
}
